import xml.etree.ElementTree as ET
import os

# 定义类别映射，这里仅关注表格，所以只有一个类别
class_mapping = {
    'table': 0
}


def xml_to_yolo_single_table(xml_file):
    """
    将单个XML标注文件转为YOLO格式的标注文件。

    :param xml_file: 输入的 XML 标注文件的路径。
    :return: list[str]: YOLO 格式的标注信息列表，每个元素是一行标注信息。
    """
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 从 XML 中获取图像的宽度和高度
    size = root.find('size')
    image_width = float(size.find('width').text)
    image_height = float(size.find('height').text)

    yolo_lines = []
    for obj in root.findall('object'):
        class_name = obj.find('name').text
        # 只处理表格类别
        if class_name == 'table':
            class_id = class_mapping[class_name]
            bbox = obj.find('bndbox')
            xmin = float(bbox.find('xmin').text)
            ymin = float(bbox.find('ymin').text)
            xmax = float(bbox.find('xmax').text)
            ymax = float(bbox.find('ymax').text)

            # 计算边界框的中心点和宽高
            x_center = (xmin + xmax) / (2 * image_width)
            y_center = (ymin + ymax) / (2 * image_height)
            width = (xmax - xmin) / image_width
            height = (ymax - ymin) / image_height

            # 生成 YOLO 格式的标注行
            yolo_line = f"{class_id} {x_center} {y_center} {width} {height}"
            yolo_lines.append(yolo_line)

    return yolo_lines


def batch_xml_to_yolo(input_dir, output_dir):
    """
    批量将输入目录中的 XML 标注文件转换为 YOLO 格式的标注文件，并保存到输出目录。

    :param input_dir: 输入的 XML 标注文件所在的目录路径。
    :param output_dir: 输出的 YOLO 格式标注文件保存的目录路径。
    :return: None
    """
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历输入目录中的所有 XML 文件
    for filename in os.listdir(input_dir):
        if filename.endswith('.xml'):
            xml_file = os.path.join(input_dir, filename)
            yolo_lines = xml_to_yolo_single_table(xml_file)

            # 生成对应的 YOLO 格式的标注文件名
            txt_filename = os.path.splitext(filename)[0] + '.txt'
            txt_file = os.path.join(output_dir, txt_filename)

            # 保存为 YOLO 格式的标注文件
            with open(txt_file, 'w') as f:
                for line in yolo_lines:
                    f.write(line + '\n')


input_dir_path = 'D:\\Edge下载\\ICDAR-2013.c-Structure\\ICDAR-2013.c-Structure\\test'  # 输入目录
output_dir_path = './ICDAR2013_output/test'
batch_xml_to_yolo(input_dir_path, output_dir_path)

input_dir_path = 'D:\\Edge下载\\ICDAR-2013.c-Structure\\ICDAR-2013.c-Structure\\val'  # 输入目录
output_dir_path = './ICDAR2013_output/val'  # 输出目录
batch_xml_to_yolo(input_dir_path, output_dir_path)
